#!/usr/bin/env python
import os,sys,time

# add self to search path for testing
if __name__ == '__main__':
    sys.path.append("../python_modules")

import general_util

import re
def fasta_sequences_iterator (lines, comment_header = "#"):
    """
        assumes file already open

        file contains either raw dna or in fasta format
        
    """
    
    comment_header = set(comment_header)
    
    curr_seq = ''
    acc = "RAW"

    line_count = 0
    for line in lines:

        line_count += 1
        

        # remove new line
        line = line.rstrip('\n')

        if not len(line):
            continue

        # ignore comment lines
        if line[0] in comment_header:
            continue
        
        # fasta header lines
        if line[0] == '>':
            if len(acc) and len(curr_seq):
                yield(acc, curr_seq)
            #elif acc != 'RAW':
            #    raise "Empty sequence on line #%d " % ( line_count - 1)
                
            curr_seq = ''

            # match identifier using regex, else use whole line
            acc = line[1:]

        # append sequence without new line
        else:
            curr_seq += line
            
    # save last sequence
    else:
        if len(curr_seq):
            yield(acc, curr_seq)
        elif acc != 'RAW':
            raise "Empty sequence on line #%d " % ( line_count - 1)


        
def unit_test():
    seq_lines = \
"""
>abc sdf
011111111110
233333333333334
55
>efg
ABBBBBBBBBBBBBBA
CDDDDDDDDDDDC
55
>hij
EF
"""
    for seq_id, seq in fasta_sequences_iterator(seq_lines.split('\n')):
        print ">"+ seq_id, "\n", seq

if __name__ == '__main__':
    unit_test()    
